import time
from contextlib import contextmanager
from datetime import datetime

from flask_sqlalchemy import SQLAlchemy as _SQLAlchemy, BaseQuery
from sqlalchemy import Column, Integer, SmallInteger

from app.libs.error_code import NotFound


class SQLAlchemy(_SQLAlchemy):
    @contextmanager
    def auto_commit(self):
        try:
            yield
            self.session.commit()
        except Exception as e:
            db.session.rollback()
            raise e


class Query(BaseQuery):
    def filter_by(self, **kwargs):
        if 'status' not in kwargs.keys():
            kwargs['status'] = 1
        return super(Query, self).filter_by(**kwargs)

    def get_or_404(self, ident, description=None):
        rv = self.get(ident)
        if rv is None:
            raise NotFound(msg='未找到')
        return rv

    def first_or_404(self, description=None):
        rv = self.first()
        if rv is None:
            raise NotFound(msg='数据未找到')
        return rv


db = SQLAlchemy(query_class=Query)


class Base(db.Model):
    __abstract__ = True
    create_time = Column(Integer, default=int(datetime.now().timestamp()))
    status = Column(SmallInteger, default=1)

    def __init__(self):
        if not self.create_time:
            self.create_time = int(datetime.now().timestamp())

    @property
    def create_datetime(self):
        if self.create_time:
            return datetime.fromtimestamp(self.create_time)
        return None

    def set_attrs(self, attrs_dict):
        for key, value in attrs_dict.items():
            if hasattr(self, key) and key != 'id':
                setattr(self, key, value)

    def add_keys(self, key):
        self.fileds.append(key)

    def delete(self):
        self.status = 0

    def get_page_data(self, model, check_data, expect_data, page, page_size):
        query = self.query.order_by(getattr(model, 'create_time').desc())
        if check_data and expect_data:
            query = query.filter(getattr(model, check_data) == expect_data).filter(model.status == 1)
        else:
            query = query.filter(model.status == 1)
        total_count = len(query.all())
        pagination = query.paginate(page=page, per_page=page_size, error_out=False)
        data_list = [data for data in pagination.items]
        return total_count, data_list

    def __getitem__(self, item):
        return getattr(self, item)
